/////////////////////////////////////////////////////////////////////////////
// Original code from libhdfs3. Copyright (c) 2013 - 2014, Pivotal Inc.
// All rights reserved. Author: Zhanwei Wang
/////////////////////////////////////////////////////////////////////////////
//  Modifications by Kumo Inc.
// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//



#include <kmhdfs/network/buffered_socket_reader.h>
#include <kmhdfs/common/exception.h>
#include <kmhdfs/common/exception_internal.h>

#include <google/protobuf/io/coded_stream.h>
#include <turbo/times/time.h>

using namespace google::protobuf::io;

namespace Hdfs {
namespace Internal {

BufferedSocketReaderImpl::BufferedSocketReaderImpl(Socket & s) :
    cursor(0), size(0), sock(s), buffer(sizeof(int64_t)) {
}

BufferedSocketReaderImpl::BufferedSocketReaderImpl(Socket & s, size_t bufferSize) :
    cursor(0), size(0), sock(s), buffer(bufferSize) {
}

int32_t BufferedSocketReaderImpl::read(char * b, int32_t s) {
    assert(s > 0 && NULL != b);
    int32_t done = s < size - cursor ? s : size - cursor;

    if (done > 0) {
        memcpy(b, buffer.data() + cursor, done);
        cursor += done;
        return done;
    } else {
        assert(size == cursor);
        size = cursor = 0;
        return sock.read(b, s);
    }
}

void BufferedSocketReaderImpl::readFully(char * b, int32_t s, int timeout) {
    assert(s > 0 && NULL != b);
    int32_t done = s < size - cursor ? s : size - cursor;
    memcpy(b, buffer.data() + cursor, done);
    cursor += done;

    if (done < s) {
        assert(size == cursor);
        size = cursor = 0;
        sock.readFully(b + done, s - done, timeout);
    }
}

int32_t BufferedSocketReaderImpl::readBigEndianInt32(int timeout) {
    char buf[sizeof(int32_t)];
    readFully(buf, sizeof(buf), timeout);
    return ntohl(*reinterpret_cast<int32_t *>(buf));
}

int32_t BufferedSocketReaderImpl::readVarint32(int timeout) {
    if (buffer.size() == 0) {
        try {
            buffer.resize(sizeof(int64_t));
            int32_t retval = readVarint32(timeout, 1);
            assert(size == cursor);
            buffer.resize(0);
            return retval;
        } catch (...) {
            assert(size == cursor);
            buffer.resize(0);
            throw;
        }
    } else {
        return readVarint32(timeout, buffer.size());
    }
}

int32_t BufferedSocketReaderImpl::readVarint32(int timeout, int32_t step) {
    int32_t value;
    bool rc = false;
    int deadline = timeout;
    memmove(&buffer[0], buffer.data() + cursor, size - cursor);
    size -= cursor;
    cursor = 0;

    while (!rc) {
        CodedInputStream in(reinterpret_cast<uint8_t *>(buffer.data() + cursor),
                            size - cursor);
        in.PushLimit(size - cursor);
        rc = in.ReadVarint32(reinterpret_cast<uint32_t *>(&value));

        if (rc) {
            cursor += size - cursor - in.BytesUntilLimit();
            return value;
        }
        turbo::TimeCost timer;
        CheckOperationCanceled();

        if (size == static_cast<int32_t>(buffer.size())) {
            THROW(HdfsNetworkException,
                  "Invalid varint type or buffer is too small, buffer size = %d.",
                  static_cast<int>(buffer.size()));
        }

        if (sock.poll(true, false, deadline)) {
            int32_t todo = buffer.size() - size;
            todo = todo < step ? todo : step;
            size += sock.read(buffer.data() + size, todo);
        }


        if (timeout > 0) {
            deadline -= timer.m_elapsed();
        }

        if (timeout >= 0 && deadline <= 0) {
            THROW(HdfsTimeoutException, "Read %d bytes timeout", size);
        }
    }

    return 0;
}

bool BufferedSocketReaderImpl::poll(int timeout) {
    if (cursor < size) {
        return true;
    }

    return sock.poll(true, false, timeout);
}

}
}
